using BotSharp.Abstraction.VectorStorage.Enums;
using BotSharp.Abstraction.VectorStorage.Options;
using System.Linq;
using Tensorflow.NumPy;

namespace BotSharp.Plugin.KnowledgeBase.MemVecDb;

public class MemoryVectorDb : IVectorDb
{
    private readonly Dictionary<string, int> _collections = new Dictionary<string, int>();
    private readonly Dictionary<string, List<VecRecord>> _vectors = new Dictionary<string, List<VecRecord>>();


    public string Provider => "MemoryVector";


    public async Task<bool> DoesCollectionExist(string collectionName)
    {
        return false;
    }

    public async Task<bool> CreateCollection(string collectionName, VectorCollectionCreateOptions options)
    {
        _collections[collectionName] = options.Dimension;
        _vectors[collectionName] = new List<VecRecord>();
        return true;
    }

    public async Task<bool> DeleteCollection(string collectionName)
    {
        return false;
    }

    public async Task<IEnumerable<string>> GetCollections()
    {
        return _collections.Select(x => x.Key).ToList();
    }

    public Task<StringIdPagedItems<VectorCollectionData>> GetPagedCollectionData(string collectionName, VectorFilter filter)
    {
        throw new NotImplementedException();
    }

    public Task<IEnumerable<VectorCollectionData>> GetCollectionData(string collectionName, IEnumerable<Guid> ids, VectorQueryOptions? options = null)
    {
        throw new NotImplementedException();
    }

    public async Task<IEnumerable<VectorCollectionData>> Search(string collectionName, float[] vector, VectorSearchOptions? options = null)
    {
        if (!_vectors.ContainsKey(collectionName))
        {
            return new List<VectorCollectionData>();
        }

        options ??= VectorSearchOptions.Default();
        var similarities = VectorHelper.CalCosineSimilarity(vector, _vectors[collectionName]);

        var results = np.argsort(similarities).ToArray<int>()
                        .Reverse()
                        .Take(options.Limit.GetValueOrDefault())
                        .Select(i => new VectorCollectionData
                        {
                            Payload = new Dictionary<string, VectorPayloadValue>
                            {
                                { "text", new(_vectors[collectionName][i].Text, VectorPayloadDataType.String) }
                            },
                            Score = similarities[i],
                            Vector = options.WithVector ? _vectors[collectionName][i].Vector : null,
                        })
                        .ToList();

        return await Task.FromResult(results);
    }

    public async Task<bool> Upsert(string collectionName, Guid id, float[] vector, string text, Dictionary<string, VectorPayloadValue>? payload = null)
    {
        _vectors[collectionName].Add(new VecRecord
        {
            Id = id.ToString(),
            Vector = vector,
            Text = text
        });

        return true;
    }

    public async Task<bool> DeleteCollectionData(string collectionName, List<Guid> ids)
    {
        return await Task.FromResult(false);
    }

    public async Task<bool> DeleteCollectionAllData(string collectionName)
    {
        return await Task.FromResult(false);
    }
}
